package com.smh.aiagent.config;

import com.knuddels.jtokkit.api.EncodingType;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import java.util.List;

/**
 * @author shiminghui
 * @date 2025/6/2 16:15
 * @description: TODO
 */
@Configuration
public class EmbeddingConfig {

    @Bean
    @Primary
    public BatchingStrategy customBatchingStrategy() {
        return new TokenCountBatchingStrategy(
                EncodingType.CL100K_BASE, // 指定编码类型
                8000, //  指定每个批次的token数量
                0.1   //  指定保留百分比
        );
    }

    @Bean
    public BatchingStrategy customBatchingStrategy2() {
        return new MyCustomBatchingStrategy();
    }

    /**
     * 自定义的 批处理策略(方法尚未实现)
     */
    public static class MyCustomBatchingStrategy implements BatchingStrategy {
        @Override
        public List<List<Document>> batch(List<Document> documents) {
            // todo
            return List.of(List.of());
        }
    }

}
